import os
import uuid
from itertools import groupby
from typing import List, Optional, Dict

import numpy as np
from common.logger import logger
from common.config import DEFAULT_CACHE_DIR
from common.typing_utils import numeric
from common.io_utils.dirs import mkdir_p
from common.io_utils.temps import temp_dir

from common.geometry.geom import Geom
from common.geometry.bbox import BoundingBox
from common.geometry.tools import project_geom
from common.geoimage.gdal_sys import mosaic_ds
from common.geoimage.raster_dataset import RasterDataset

from catalog.models.base_stac_item import BaseSTACItem
from catalog.raster.utils import (
    determine_dest_srs,
    determine_dest_res,
    determine_dest_extent,
)
from catalog.raster.utils import (
    is_s2_l2a_collection,
    update_s2_l2a_boa,
    get_s2_l2a_add_boa_offset,
    nb_customized_median,
)
from enhance.cloudseg.cloud_seg_service import CloudSegService

TMP_DIR = os.path.join(DEFAULT_CACHE_DIR, "catalog", "raster")
mkdir_p(TMP_DIR)


def cloud_mask_per_mgrs_tile(
    temp_dir_path: str, item: BaseSTACItem, aoi: Geom, cloud_seg: CloudSegService
):
    """Extract cloud mask on the tile with Cloud Segmentation API, return a binary mask of cloud region.

    Parameters
    ----------
    temp_dir_path : str
        A folder to store the temporary processing results of Cloud Segmentation API.
    item : BaseSTACItem
        The item of Sentinel-2 product.
    aoi : Geom
        The AOI of item
    cloud_seg : CloudSegService
        The class instance of CloudSegService

    Returns
    -------
    RasterDataset or None
        Return the binary map of cloud region in RasterDataset format, or
        Return None for encountering Exception.
    """
    # NOTE Cloud API takes RGBN as input
    asset_names = ["B04", "B03", "B02", "B08"]
    # NOTE We skip the resolution warp because the RGBN has the same resolution.
    vrt_scene = item.choose_assets(asset_names)
    rds = RasterDataset.from_vrt_scene(vrt_scene, aoi=aoi)

    # NOTE Post-process Sentinel-2 L2A due to ESA updates of -1000 offset
    asset_offsets = get_s2_l2a_add_boa_offset(item, asset_names)
    if asset_offsets is not None:
        rds.data = update_s2_l2a_boa(rds.data, asset_offsets)

    try:
        rds_cloud_seg = cloud_seg.inference_raster(rds)
    except Exception as err:
        msg = f"[Error] Failed to run S2 cloud segmentation due to `{err}`."
        logger.error(msg)
        print(msg)
        return None

    # Class 0: Clear Ground,
    # Class 1: Cloud Shadow
    # Class 2: Cirrus (thin cloud)
    # Class 3: Cloud (thick cloud)
    # Class 4: No-Data
    cloud_mask = rds_cloud_seg.data > 0
    cloud_mask = np.squeeze(cloud_mask)
    return cloud_mask


def composite_s2decloud_v2_per_mgrs_tile(
    items: List[BaseSTACItem],
    asset_names: Dict[str, str],
    aoi: Geom,
    temp_dir_path: str,
    cloud_seg: CloudSegService,
):
    """composite a list of items in the same MGRS tile, using s2decloud-v2 method
    based on Cloud Segmentation API removing cloud, cirrus, cloud shadow, and nodata.

    Parameters
    ----------
    item : BaseSTACItem
        The item of Sentinel-2 product.
    asset_names : Dict[str, str]
        Asset Names of Sentinel-2 item, e.g., ["B02", "B03", "B04"]
    aoi : Geom
        The AOI of item
    temp_dir_path : str
        A folder to store the temporary processing results of Cloud Segmentation API.
    cloud_seg : CloudSegService
        The class instance of CloudSegService

    Returns
    -------
    RasterDataset
        Return the binary map of cloud region in RasterDataset format
    """

    if len(items) == 1:
        item = items[0]
        vrt_scene = item.choose_assets(asset_names)
        rds = RasterDataset.from_vrt_scene(vrt_scene, aoi=aoi)
        # NOTE Post-process Sentinel-2 L2A due to ESA updates of -1000 offset
        asset_offsets = get_s2_l2a_add_boa_offset(item, asset_names)
        if asset_offsets is not None:
            rds.data = update_s2_l2a_boa(rds.data, asset_offsets)

        # NOTE Call the Could Segmentation API
        cloud_mask = cloud_mask_per_mgrs_tile(temp_dir_path, item, aoi, cloud_seg)

        meta = rds.meta
        nda_data = []
        for band_ind in range(meta.n_bands):
            nda_data.append(np.where(cloud_mask, 0, rds.data[band_ind, :, :]))
        nda_data = np.stack(nda_data, axis=0)
        return RasterDataset.from_ndarray(nda_data, meta)

    loaded_rds = []
    for item in items:
        vrt_scene = item.choose_assets(asset_names)
        rds = RasterDataset.from_vrt_scene(vrt_scene, aoi=aoi)
        # NOTE Post-process Sentinel-2 L2A due to ESA updates of -1000 offset
        asset_offsets = get_s2_l2a_add_boa_offset(item, asset_names)
        if asset_offsets is not None:
            rds.data = update_s2_l2a_boa(rds.data, asset_offsets)
        loaded_rds.append(rds)

    stacked_cloud_mask = []
    for item in items:
        # NOTE Call the Could Segmentation API
        cloud_mask = cloud_mask_per_mgrs_tile(temp_dir_path, item, aoi, cloud_seg)
        stacked_cloud_mask.append(cloud_mask)
    stacked_cloud_mask = np.stack(stacked_cloud_mask).astype(np.uint8)

    nda_res = []
    for ind_band, _ in enumerate(asset_names):
        nda_data_tmp = np.stack(
            [np.squeeze(rds_tmp.data[ind_band, :, :]) for rds_tmp in loaded_rds]
        )
        nda_res.append(nb_customized_median(nda_data_tmp, stacked_cloud_mask))
    nda_res = np.stack(nda_res)

    return RasterDataset.from_ndarray(nda_res, loaded_rds[-1].meta)


def composite_s2decloud_v2(
    items: List[BaseSTACItem],
    asset_names: List[str],
    extent: Optional[BoundingBox] = None,
    aoi: Optional[Geom] = None,
    dest_srs: Optional[str] = None,
    dest_res: Optional[numeric] = None,
    resample_alg: str = "nearest",
):
    """Composite a given list of `BaseSTACItems` using a median composite method,
    based on Cloud Segmentation API removing cloud, cirrus, cloud shadow, and nodata.

    Parameters
    ----------
    items : List[BaseSTACItem]
        A list of BaseSTACItems to composite
    asset_names: List[str]
        A list of asset names, e.g. ["B02", "B03", "B04"]
    extent: Optional[BoundingBox], optional
        The extent of AOI in BoundingBox. The output will be cropped to the shape if given.
    aoi: Optional[Geom], optional
        The AOI in Geom format. The output will be cropped to the shape if given.
    dest_srs: Optional[str], optional
        The SRS of the destination raster. The function will computer the centroid of AOI if not given.
    dest_res: Optional[numeric], optional
        The pixel ground distance of the destination ratser. Default is 10-meter as Sentinel-2 resolution.
    resample_alg: string, optional
        The resampling algorithm of image warping. One of the options
        "nearest", "bilinear", "cubic", "cubicspline", "lanczos", "average", and "mode". Default is "nearest".

    Returns
    -------
    RasterDataset
        A composited result in RasterDataset
    """

    if not is_s2_l2a_collection(items):
        raise ValueError(
            (
                "The input items are not from the COPERNICUS/SENTINEL2/L2A product"
                " and not supported by this composite strategy."
            )
        )

    if dest_srs is None:
        dest_srs = determine_dest_srs(aoi, items)

    if dest_res is None:
        dest_res = determine_dest_res(dest_srs, items)

    dest_extent = determine_dest_extent(extent, aoi, dest_srs, dest_res)
    if dest_extent is not None:
        # if an extent is given, convert it to a geometry under WGS84 as the AOI
        bbox_geom = Geom.from_bbox(dest_extent)
        if dest_srs.upper() != "EPSG:4326":
            work_aoi = project_geom(bbox_geom, src_epsg=dest_srs, dest_epsg="EPSG:4326")
        else:
            work_aoi = bbox_geom
    else:
        work_aoi = None

    # group by mgrs tile
    items = sorted(items, key=lambda item: item.mgrs_tile)
    tile_groups = groupby(items, key=lambda item: item.mgrs_tile)
    n_tiles = len(set(item.mgrs_tile for item in items))

    cloud_seg = CloudSegService()
    temp_dir_name = str(uuid.uuid4())
    if n_tiles == 1:
        with temp_dir(
            dir_name=temp_dir_name, base_dir=TMP_DIR, remove=True
        ) as temp_dir_path:
            for mgrs_tile, sub_items in tile_groups:
                sub_items = list(sub_items)
            logger.info(f"working on {mgrs_tile} with {len(sub_items)} images")
            rds = composite_s2decloud_v2_per_mgrs_tile(
                sub_items, asset_names, work_aoi, temp_dir_path, cloud_seg
            )
            rds = rds.warp(
                dest_srs=dest_srs, dest_res=dest_res, dest_extent=dest_extent
            )
        return rds

    per_tile_files = []
    with temp_dir(
        dir_name=temp_dir_name, base_dir=TMP_DIR, remove=True
    ) as temp_dir_path:
        print(f"All tile_groups: {tile_groups}")
        for mgrs_tile, sub_items in tile_groups:
            sub_items = list(sub_items)
            logger.info(f"working on {mgrs_tile} with {len(sub_items)} images")
            rds = composite_s2decloud_v2_per_mgrs_tile(
                sub_items, asset_names, work_aoi, temp_dir_path, cloud_seg
            )
            if dest_srs != rds.meta.prj_epsg:
                rds = rds.warp(dest_srs=dest_srs)
            fpath_tmp = os.path.join(temp_dir_path, f"{mgrs_tile}.tif")
            rds.to_geotiff(fpath_tmp)
            per_tile_files.append(fpath_tmp)

        with mosaic_ds(
            per_tile_files,
            resample_alg=resample_alg,
            dest_res=dest_res,
            dest_extent=dest_extent,
            src_nodata=0,
        ) as gdal_rst:
            rds = RasterDataset.from_gdal_raster(gdal_rst, extent=dest_extent)

    return rds
